package org.example;

import java.io.*;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

public class PassToMd5Main {

    private static final int BATCH_SIZE = 10000;
    private static final int IO_BUFFER_SIZE = 1 << 20; // 1MB 缓冲
    private static final char[] HEX_ARRAY = "0123456789abcdef".toCharArray();
    private static final ThreadLocal<MessageDigest> THREAD_LOCAL_MD5 =
            ThreadLocal.withInitial(() -> {
                try {
                    return MessageDigest.getInstance("MD5");
                } catch (NoSuchAlgorithmException e) {
                    throw new RuntimeException("MD5 not available", e);
                }
            });

    public static void processLargeFile(String filePath, String outputPath) throws Exception {
        AtomicLong lineCount = new AtomicLong(0);
        long startTime = System.currentTimeMillis();

        int cpuThreads = Runtime.getRuntime().availableProcessors();
        ExecutorService executor = Executors.newFixedThreadPool(cpuThreads);
        Semaphore semaphore = new Semaphore(cpuThreads * 2); // 限制并行批次数

        try (BufferedReader reader = new BufferedReader(
                new FileReader(filePath, StandardCharsets.UTF_8), IO_BUFFER_SIZE);
             BufferedWriter writer = new BufferedWriter(
                     new OutputStreamWriter(new FileOutputStream(outputPath, true), StandardCharsets.UTF_8),
                     IO_BUFFER_SIZE)) {

            List<String> batch = new java.util.ArrayList<>(BATCH_SIZE);
            String line;
            while ((line = reader.readLine()) != null) {
                batch.add(line);
                lineCount.incrementAndGet();

                if (batch.size() >= BATCH_SIZE) {
                    List<String> submitBatch = batch;
                    batch = new java.util.ArrayList<>(BATCH_SIZE);

                    semaphore.acquire(); // 控制并行批次数
                    executor.submit(() -> {
                        try {
                            processAndWrite(submitBatch, writer);
                        } finally {
                            semaphore.release();
                        }
                    });
                }

                if (lineCount.get() % 1_000_000 == 0) {
                    reportProgress(lineCount.get(), startTime);
                }
            }

            if (!batch.isEmpty()) {
                semaphore.acquire();
                List<String> submitBatch = batch;
                executor.submit(() -> {
                    try {
                        processAndWrite(submitBatch, writer);
                    } finally {
                        semaphore.release();
                    }
                });
            }
        }

        executor.shutdown();
        executor.awaitTermination(1, TimeUnit.HOURS);

        long endTime = System.currentTimeMillis();
        double totalSeconds = (endTime - startTime) / 1000.0;

        System.out.printf("Finished: %,d lines in %.1f sec (%.1f lines/sec)%n",
                lineCount.get(), totalSeconds, lineCount.get() / totalSeconds);
        System.out.printf("Total elapsed time: %.3f seconds%n", totalSeconds);
        // ✅ Finished writing 320,145,783 lines in 1131.50 seconds (282939.77 lines/sec) 18分钟
    }

    private static void processAndWrite(List<String> batch, BufferedWriter writer) {
        try {
            StringBuilder sb = new StringBuilder(batch.size() * 64);
            for (String line : batch) {
                String md5 = calculateMD5(line);
                sb.append("{\"text\":\"").append(escapeJson(line)).append("\",\"md5\":\"").append(md5).append("\"}\n");
                if (sb.length() > 64 * 1000) { // 每 1000 条写一次
                    synchronized (writer) {
                        writer.write(sb.toString());
                    }
                    sb.setLength(0);
                }
            }
            if (sb.length() > 0) {
                synchronized (writer) {
                    writer.write(sb.toString());
                }
            }
        } catch (IOException e) {
            throw new UncheckedIOException(e);
        }
    }

    private static String calculateMD5(String input) {
        MessageDigest md5 = THREAD_LOCAL_MD5.get();
        md5.reset();
        byte[] digest = md5.digest(input.getBytes(StandardCharsets.UTF_8));
        char[] hexChars = new char[32];
        for (int i = 0; i < 16; i++) {
            int v = digest[i] & 0xFF;
            hexChars[i * 2] = HEX_ARRAY[v >>> 4];
            hexChars[i * 2 + 1] = HEX_ARRAY[v & 0x0F];
        }
        return new String(hexChars);
    }

    private static String escapeJson(String text) {
        StringBuilder sb = new StringBuilder(text.length() + 16);
        for (char c : text.toCharArray()) {
            switch (c) {
                case '\\' -> sb.append("\\\\");
                case '"' -> sb.append("\\\"");
                case '\n' -> sb.append("\\n");
                case '\r' -> sb.append("\\r");
                case '\t' -> sb.append("\\t");
                default -> sb.append(c);
            }
        }
        return sb.toString();
    }

    private static void reportProgress(long lineCount, long startTime) {
        long currentTime = System.currentTimeMillis();
        double elapsedSeconds = (currentTime - startTime) / 1000.0;
        double linesPerSecond = lineCount / elapsedSeconds;
        System.out.printf("Progress: %,d lines (%.1f lines/sec, %.1f min elapsed)%n",
                lineCount, linesPerSecond, elapsedSeconds / 60);
    }

    public static void main(String[] args) throws Exception {
        String filePath = "weakpass_4.policy.txt";
        String outputPath = "md5.jsonl";
        System.out.println("Processing file: " + filePath);
        processLargeFile(filePath, outputPath);
        // Total elapsed time: 121.551 seconds
    }
}
